
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
import torchvision
from utils.dataloader import DataGenerator
from utils.utils import show_config



def inverse_preprocess_input(x):
    x = (x + 1.) / 2. * 255
    return x

#-------------------------------------#
#   设置数据集的路径
#   train_txt_path   训练图片路径和标签
#   valid_txt_path   验证图片路径和标签
#-------------------------------------#
train_txt_path   = "cls_train.txt"
valid_txt_path   = "cls_valid.txt"
input_shape = [300, 300]

with open(train_txt_path, "r") as f:
    trainlines = f.readlines()

with open(valid_txt_path, "r") as f:
    validlines = f.readlines()

# 数据读取
# 构建DataGenerator实例

train_data = DataGenerator(trainlines, input_shape=input_shape, random=True)
valid_data = DataGenerator(validlines, input_shape=input_shape, random=False)

# 构建DataLoder
train_loader = DataLoader(dataset=train_data, batch_size=24, shuffle=True)
valid_loader = DataLoader(dataset=valid_data, batch_size=8)


if __name__ == "__main__":

    # 获得数据集的长度 len(), 即length
    train_data_size = len(train_data)
    valid_data_size = len(valid_data)

    # 格式化字符串, format() 中的数据会替换 {}
    print("训练数据集及的长度为: {}".format(train_data_size))
    print("验证数据集及的长度为: {}".format(valid_data_size))

    show_config(train_data_num=train_data_size, valid_data_num=valid_data_size)

    for iteration, (batch_x, batch_y) in enumerate(train_loader):
        # if iteration < 1:
        a = batch_x
        b = batch_y
        image_batch = inverse_preprocess_input(batch_x)
        image_batch = torchvision.utils.make_grid(image_batch, padding=0)
        image_batch = image_batch.numpy().astype(np.uint8)
        print(image_batch.shape)
        imgs = np.transpose(image_batch, (1, 2, 0))
        print(imgs.shape)
        plt.figure(figsize=(12, 8))
        plt.imshow(imgs)
        plt.tight_layout()
        plt.show()
